import json
import os
from typing import Dict, Any, Optional

from preprocess.name_sep import extract_name_from_text
from preprocess.seperate_key_part import extract_key_sections, extract_docx_text

from pipline.Admission_Record import extract_admission_diagnosis, convert_to_diagnosis_dict
from pipline.Death_Record import extract_diagnosis_info_enhanced
from pipline.Death_Case_Discussion import extract_admission_diagnosis_and_history


def main(doc_path: str, output_dir: Optional[str] = None) -> Dict[str, Any]:
    """
    处理医疗文档（.docx），提取姓名、入院记录、死亡记录、死亡病例讨论等关键信息。

    参数:
        doc_path (str): 医疗文档的完整路径（.docx 文件）。
        output_dir (str, optional): 用于保存中间文本文件的目录。若为 None，则使用文档所在目录。

    返回:
        dict: 包含提取结果的字典，键包括：
            - 'name': 患者姓名
            - 'admission_diagnosis': 入院诊断（结构化字典）
            - 'death_record': 死亡记录中的诊断与原因
            - 'death_discussion': 死亡病例讨论中的入院诊断与病史摘要
    """
    if not os.path.isfile(doc_path):
        raise FileNotFoundError(f"文档不存在: {doc_path}")

    # 默认输出目录为文档所在目录
    if output_dir is None:
        output_dir = os.path.dirname(doc_path)

    # 提取纯文本
    text = extract_docx_text(doc_path)

    # 提取姓名
    name = extract_name_from_text(text)
    if not name:
        name = "未知姓名"

    # 构造中间文本文件名
    base_name = os.path.splitext(os.path.basename(doc_path))[0]
    txt_filename = f"{base_name}_{name}.txt"
    txt_path = os.path.join(output_dir, txt_filename)

    # 提取关键章节（会生成 txt 文件）
    sections = extract_key_sections(doc_path, txt_path)

    # 初始化结果字典
    result = {
        'name': name,
        'admission_diagnosis': None,
        'death_record': None,
        'death_discussion': None
    }

    # 提取入院记录中的入院诊断
    try:
        if '入院记录' in sections and sections['入院记录']:
            raw_admission = extract_admission_diagnosis(sections['入院记录'])
            result['admission_diagnosis'] = convert_to_diagnosis_dict(raw_admission)
    except Exception as e:
        print(f"处理入院记录时出错: {e}")

    # 提取死亡记录中的诊断信息
    try:
        if '死亡记录' in sections and sections['死亡记录']:
            result['death_record'] = extract_diagnosis_info_enhanced(sections['死亡记录'])
    except Exception as e:
        print(f"处理死亡记录时出错: {e}")

    # 提取死亡病例讨论中的信息
    try:
        if '死亡病例讨论' in sections and sections['死亡病例讨论']:
            result['death_discussion'] = extract_admission_diagnosis_and_history(sections['死亡病例讨论'])
    except Exception as e:
        print(f"处理死亡病例讨论时出错: {e}")

    # === 新增：将 result 保存为 {base_name}_{name}_seperate.txt ===
    result_filename = f"{base_name}_{name}_seperate.txt"
    result_path = os.path.join(output_dir, result_filename)

    # 使用 json.dumps 美化输出，确保中文可读
    with open(result_path, 'w', encoding='utf-8') as f:
        json.dump(result, f, ensure_ascii=False, indent=4)

    print(f"结果已保存至: {result_path}")

    return result


# 示例调用
if __name__ == "__main__":
    dict_path = r'C:\Users\yyds\Desktop\new_demo'
    doc_file = os.path.join(dict_path, "519038.docx")

    try:
        results = main(doc_file)
        print("患者姓名:", results['name'])
        print("入院诊断:", results['admission_diagnosis'])
        print("死亡记录:", results['death_record'])
        print("死亡病例讨论:", results['death_discussion'])
    except Exception as e:
        print(f"处理文档时发生错误: {e}")